'''
https://leetcode.cn/problems/optimize-water-distribution-in-a-village
'''
from typing import List

class UF:
    def __init__(self, n):
        self.father = [i for i in range(n)]
    def find(self, x):
        if x != self.father[x]:
            self.father[x] = self.find(self.father[x])
        return self.father[x]
    def union(self, x, y):
        fx, fy = self.find(x), self.find(y)
        if fx == fy: return False
        self.father[fy] = fx
        return True

class Solution:

    # 建立虚拟节点，将每个节点自身建水井的消耗看作与虚拟节点连接的边
    #       然后转化问题为，将所有节点，与虚拟节点，连接的最小生成树
    # 最小生成树
    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:
        virtual_node = 0
        for i, well in enumerate(wells):
            u = i + 1
            pipes.append([virtual_node, u, well])

        pipes.sort(key=lambda x: x[2])
        uf = UF(n + 1)
        n_conn_edges = 0
        total_cost = 0
        for u, v, w in pipes:
            if uf.union(u, v):
                n_conn_edges += 1
                total_cost += w
            if n_conn_edges == n:
                return total_cost
        return -1

n = 3
wells = [1,2,2]
pipes = [[1,2,1],[2,3,1]]
print(Solution().minCostToSupplyWater(n, wells, pipes))
